/*
 * Copyright (C) 2007 Eskil Bylund
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Security.Cryptography;
using System.Text;
using System.Xml;

namespace DCSharp.Hashing
{
	/// <summary>
	/// A class to store hash trees.
	/// </summary>
	public partial class HashStore
	{
		private Dictionary<string, HashFileInfo> fileInfos;
		private Dictionary<string, HashInfo> hashInfos;
		private HashAlgorithm algorithm;

		private string indexFile;
		private string dataFile;

		public HashStore(HashAlgorithm algorithm, string indexFile,
			string dataFile)
		{
			this.algorithm = algorithm;
			this.indexFile = indexFile;
			this.dataFile = dataFile;

			// TODO: Remove OrdinalIgnoreCase when the filename is no longer lowercased. 
			fileInfos = new Dictionary<string, HashFileInfo>(StringComparer.OrdinalIgnoreCase);
			hashInfos = new Dictionary<string, HashInfo>();
		}

		#region Properties

		/// <summary>
		/// Gets the index file.
		/// </summary>
		/// <value>The path to the index file.</value>
		public string IndexFile
		{
			get { return indexFile; }
		}

		/// <summary>
		/// Gets the data file.
		/// </summary>
		/// <value>The path to the data file.</value>
		public string DataFile
		{
			get { return dataFile; }
		}

		/// <summary>
		/// Gets the number of stored files.
		/// </summary>
		/// <value>The number of stored files.</value>
		public int FileCount
		{
			get { return fileInfos.Count; }
		}

		/// <summary>
		/// Gets the number of stored hash trees.
		/// </summary>
		/// <value>The number of stored hash trees.</value>
		public int HashTreeCount
		{
			get { return hashInfos.Count; }
		}

		#endregion

		#region Methods

		/// <summary>
		/// Adds the hash tree for a file to the store. 
		/// </summary>
		/// <param name="fileInfo">The file that the hash tree was created from.</param>
		/// <param name="hashTree">The hash tree to add.</param>
		public void AddFile(FileInfo fileInfo, HashTree hashTree)
		{
			if (fileInfo == null)
			{
				throw new ArgumentNullException("fileInfo");
			}
			AddHashTree(hashTree);

			string filename = fileInfo.FullName;

			// Remove old HashFileInfo since this might be an update
			fileInfos.Remove(filename);

			HashFileInfo hashFileInfo = new HashFileInfo(filename,
				ToSecondsSinceEpoch(fileInfo.LastWriteTimeUtc),
				Base32.Encode(hashTree.Root));

			fileInfos.Add(filename, hashFileInfo);
		}

		/// <summary>
		/// Adds a hash tree to the store.
		/// </summary>
		/// <param name="hashTree">The hash tree to add.</param>
		public void AddHashTree(HashTree hashTree)
		{
			if (hashTree == null)
			{
				throw new ArgumentNullException("hashTree");
			}

			string root = Base32.Encode(hashTree.Root);
			if (hashInfos.ContainsKey(root))
			{
				return;
			}

			// TODO: Let the caller handle the exception?
			try
			{
				HashInfo hashInfo = WriteHashTree(dataFile, hashTree);
				hashInfos.Add(root, hashInfo);
			}
			catch (Exception e)
			{
				Debug.WriteLine("Unable to write HashTree: " + e);
			}
		}

		/// <summary>
		/// Checks if the hash tree for a file exists in the store.
		/// </summary>
		/// <param name="fileInfo">The file to check.</param>
		/// <return>True if the HashTree exists; otherwise, false.</return>
		public bool HasHash(FileInfo fileInfo)
		{
			if (fileInfo == null)
			{
				throw new ArgumentNullException("fileInfo");
			}

			string filename = fileInfo.FullName;

			HashFileInfo hashFileInfo;
			HashInfo hashInfo;

			if (fileInfos.TryGetValue(filename, out hashFileInfo) &&
				hashInfos.TryGetValue(hashFileInfo.Root, out hashInfo))
			{
				hashFileInfo.KeepOnRebuild = true;

				// Convert to seconds since the stored date does not contains the exact milliseconds.
				int lastWriteTime = ToSecondsSinceEpoch(fileInfo.LastWriteTimeUtc);
				if (hashInfo.Size == fileInfo.Length &&
					hashFileInfo.LastWriteTime == lastWriteTime)
				{
					return true;
				}
			}
			return false;
		}

		/// <summary>
		/// Returns the root hash for a file.
		/// </summary>
		/// <param name="fileInfo">The file to get the root hash for.</param>
		/// <returns>
		/// The Base32 encoded root hash, or null if the hash wasn't found.
		/// </returns>
		public string GetRootHash(FileInfo fileInfo)
		{
			if (fileInfo == null)
			{
				throw new ArgumentNullException("fileInfo");
			}

			string filename = fileInfo.FullName;

			HashFileInfo hashFileInfo;
			if (fileInfos.TryGetValue(filename, out hashFileInfo))
			{
				hashFileInfo.KeepOnRebuild = true;
				return hashFileInfo.Root;
			}
			return null;
		}

		/// <summary>
		/// Gets the hash tree with the specified root hash.
		/// </summary>
		/// <param name="root">The Base32 encoded root hash.</param>
		/// <returns>The HashTree if found; otherwise, null.</returns>
		public HashTree GetHashTree(string root)
		{
			if (root == null)
			{
				throw new ArgumentNullException("root");
			}

			HashInfo hashInfo;
			if (hashInfos.TryGetValue(root, out hashInfo))
			{
				// TODO: Let the caller handle the exception?
				try
				{
					return ReadHashTree(dataFile, hashInfo);
				}
				catch (Exception e)
				{
					Debug.WriteLine("Unable to read HashTree: " + e);
				}
			}
			return null;
		}

		/// <summary>
		/// Removes any obsolete files and hash trees and rebuilds the data file.
		/// </summary>
		public void Rebuild()
		{
			Dictionary<string, HashInfo> tempHashInfos = new Dictionary<string, HashInfo>();

			string tempFile = dataFile + ".tmp";
			File.Delete(tempFile);

			string[] filenames = new string[fileInfos.Keys.Count];
			fileInfos.Keys.CopyTo(filenames, 0);

			foreach (string filename in filenames)
			{
				HashFileInfo fileInfo;
				HashInfo hashInfo = null;

				if (fileInfos.TryGetValue(filename, out fileInfo) &&
					fileInfo.Root != null && fileInfo.KeepOnRebuild)
				{
					if (tempHashInfos.ContainsKey(fileInfo.Root))
					{
						continue;
					}

					HashTree hashTree = GetHashTree(fileInfo.Root);
					if (hashTree != null)
					{
						hashInfo = WriteHashTree(tempFile, hashTree);
						tempHashInfos.Add(hashInfo.Root, hashInfo);
					}
				}

				if (hashInfo == null)
				{
					// No hash tree written
					fileInfos.Remove(filename);
				}
			}

			hashInfos = tempHashInfos;

			File.Delete(dataFile);
			if (File.Exists(tempFile))
			{
				File.Move(tempFile, dataFile);
			}
		}

		#region IndexFile

		/// <summary>
		/// Saves the store to disk.
		/// </summary>
		public void Save()
		{
			// Omit the byte order mark
			Encoding encoding = new UTF8Encoding(false);

			using (XmlTextWriter writer = new XmlTextWriter(indexFile, encoding))
			{
				writer.Formatting = Formatting.Indented;
				writer.IndentChar = '\t';
				writer.Indentation = 1;

				Save(writer);
			}
		}

		protected void Save(XmlTextWriter writer)
		{
			if (writer == null)
			{
				throw new ArgumentNullException("writer");
			}

			// Include the standalone attribute
			writer.WriteStartDocument(true);

			writer.WriteStartElement("HashStore");
			writer.WriteAttributeString("Version", "2");

			if (hashInfos.Count > 0)
			{
				writer.WriteStartElement("Trees");
				foreach (HashInfo info in hashInfos.Values)
				{
					writer.WriteStartElement("Hash");
					writer.WriteAttributeString("Type", info.Type);
					writer.WriteAttributeString("Index", info.Index.ToString());
					writer.WriteAttributeString("BlockSize", info.BlockSize.ToString());
					writer.WriteAttributeString("Size", info.Size.ToString());
					writer.WriteAttributeString("Root", info.Root);
					writer.WriteEndElement();
				}
				writer.WriteEndElement();
			}
			if (fileInfos.Count > 0)
			{
				writer.WriteStartElement("Files");
				foreach (HashFileInfo fileInfo in fileInfos.Values)
				{
					// FIXME: Lowercasing the file name is stupid, but that's how DC++ does it
					writer.WriteStartElement("File");
					writer.WriteAttributeString("Name", fileInfo.Name.ToLower());
					writer.WriteAttributeString("TimeStamp", fileInfo.LastWriteTime.ToString());
					writer.WriteAttributeString("Root", fileInfo.Root);
					writer.WriteEndElement();
				}
				writer.WriteEndElement();
			}

			writer.WriteEndElement();
		}

		/// <summary>
		/// Loads the store from disk.
		/// </summary>
		public void Load()
		{
			hashInfos.Clear();
			fileInfos.Clear();

			using (XmlTextReader reader = new XmlTextReader(indexFile))
			{
				Load(reader);
			}
		}

		protected void Load(XmlTextReader reader)
		{
			if (reader == null)
			{
				throw new ArgumentNullException("reader");
			}
			if (reader.ReadToFollowing("Trees"))
			{
				XmlReader subtree = reader.ReadSubtree();
				while (subtree.ReadToFollowing("Hash"))
				{
					try
					{
						HashInfo hashInfo = new HashInfo();

						hashInfo.Type = subtree.GetAttribute("Type");
						hashInfo.Index = long.Parse(subtree.GetAttribute("Index"));
						hashInfo.BlockSize = long.Parse(subtree.GetAttribute("BlockSize"));
						hashInfo.Size = long.Parse(subtree.GetAttribute("Size"));
						hashInfo.Root = subtree.GetAttribute("Root");

						if (IsValid(hashInfo))
						{
							hashInfos.Add(hashInfo.Root, hashInfo);
						}
					}
					catch
					{
						// Ignore invalid hash info
					}
				}
			}
			if (reader.ReadToFollowing("Files"))
			{
				XmlReader subtree = reader.ReadSubtree();
				while (subtree.ReadToFollowing("File"))
				{
					HashFileInfo fileInfo = new HashFileInfo();
					try
					{
						fileInfo.Name = subtree.GetAttribute("Name");
						fileInfo.LastWriteTime = int.Parse(subtree.GetAttribute("TimeStamp"));
						fileInfo.Root = subtree.GetAttribute("Root");
						fileInfo.KeepOnRebuild = false;

						fileInfos.Add(fileInfo.Name, fileInfo);
					}
					catch
					{
						// Ignore invalid file info
					}
				}
			}
		}

		#endregion

		#region DataFile

		// TODO
		protected HashAlgorithm GetAlgorithm(string type)
		{
			return algorithm;
		}

		protected HashTree ReadHashTree(string dataFile, HashInfo hashInfo)
		{
			if (hashInfo.Root == null)
			{
				throw new ArgumentException("Missing root in HashInfo",
					"hashInfo");
			}

			byte[] root = Base32.Decode(hashInfo.Root);

			if (hashInfo.Index == -1)
			{
				Stream stream = new MemoryStream(root);

				return new HashTree(GetAlgorithm(hashInfo.Type),
					hashInfo.BlockSize, hashInfo.Size, stream, root);
			}

			using (Stream stream = new FileStream(dataFile, FileMode.Open,
				FileAccess.Read, FileShare.ReadWrite))
			{
				stream.Seek(hashInfo.Index, SeekOrigin.Begin);

				return new HashTree(GetAlgorithm(hashInfo.Type),
					hashInfo.BlockSize, hashInfo.Size, stream, root);
			}
		}

		protected HashInfo WriteHashTree(string dataFile, HashTree hashTree)
		{
			HashInfo hashInfo = new HashInfo(0, hashTree.BlockSize,
				hashTree.FileSize, Base32.Encode(hashTree.Root));

			if (hashTree.LeafCount == 1)
			{
				hashInfo.Index = -1;
				return hashInfo;
			}

			using (FileStream stream = new FileStream(dataFile,
				FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.Read))
			{
				byte[] buffer = new byte[8];
				long end = buffer.Length;

				// Get the end position
				if (stream.Read(buffer, 0, buffer.Length) == buffer.Length)
				{
					end = BitConverter.ToInt64(buffer, 0);
				}
				hashInfo.Index = end;

				// Grow the file?
				if (end + hashTree.LeafCount * algorithm.HashSize > stream.Length)
				{
					stream.Seek(stream.Length, SeekOrigin.Begin);
					for (int i = 0; i < 1024 * 1024; i++)
					{
						stream.WriteByte(0);
					}
				}

				// Write the leaves
				stream.Seek(end, SeekOrigin.Begin);
				foreach (byte[] leaf in hashTree.Leaves)
				{
					stream.Write(leaf, 0, leaf.Length);
				}

				// Update the end position
				end = stream.Position;
				buffer = BitConverter.GetBytes(end);

				stream.Seek(0, SeekOrigin.Begin);
				stream.Write(buffer, 0, buffer.Length);
			}
			return hashInfo;
		}

		#endregion

		// TODO: Switch to long in 2038...
		private static int ToSecondsSinceEpoch(DateTime date)
		{
			return (int)date.Subtract(new DateTime(1970, 1, 1)).TotalSeconds;
		}

		private static bool IsValid(HashInfo hashInfo)
		{
			return hashInfo.Root != null && hashInfo.Size >= 0 &&
				hashInfo.BlockSize >= HashTree.SegmentSize;
		}

		#endregion
	}
}
